import cv2
import numpy as np
import os
from PIL import Image, ImageEnhance, ImageFilter, ImageOps


def enhance_image_clarity(image_path):
    """增强图片清晰度，特别针对白色/蓝色底色的用例图优化文字清晰度"""
    try:
        # 读取图片并转换为PIL格式
        original_img = Image.open(image_path)

        # 转换为OpenCV格式 (BGR)
        cv_img = np.array(original_img)
        if cv_img.ndim == 3 and cv_img.shape[2] == 4:
            cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGR)
            has_alpha = True
        else:
            has_alpha = False

        # 检测图片主要背景色（白色或蓝色）
        hsv_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2HSV)

        # 定义白色和蓝色的HSV范围
        lower_white = np.array([0, 0, 200], dtype=np.uint8)
        upper_white = np.array([180, 30, 255], dtype=np.uint8)

        lower_blue = np.array([90, 50, 50], dtype=np.uint8)
        upper_blue = np.array([130, 255, 255], dtype=np.uint8)

        # 创建掩码识别白色和蓝色区域
        white_mask = cv2.inRange(hsv_img, lower_white, upper_white)
        blue_mask = cv2.inRange(hsv_img, lower_blue, upper_blue)

        # 组合背景掩码
        background_mask = cv2.bitwise_or(white_mask, blue_mask)

        # 反转掩码获取文字区域
        text_mask = cv2.bitwise_not(background_mask)

        # 增强文字对比度 - 针对文字区域特别处理
        # 方法1：自适应直方图均衡化
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        enhanced_gray = clahe.apply(gray)

        # 方法2：非锐化掩蔽增强边缘
        blurred = cv2.GaussianBlur(enhanced_gray, (0, 0), 3)
        unsharp_mask = cv2.addWeighted(enhanced_gray, 1.5, blurred, -0.5, 0)

        # 方法3：二值化处理优化文字
        _, binary = cv2.threshold(unsharp_mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # 创建三通道文字增强图像
        text_enhanced = cv2.merge([binary, binary, binary])

        # 融合处理结果：背景保留原色，文字使用增强后的效果
        # 转换掩码为三通道
        background_mask_3ch = cv2.merge([background_mask, background_mask, background_mask])
        text_mask_3ch = cv2.merge([text_mask, text_mask, text_mask])

        # 提取背景
        background = cv2.bitwise_and(cv_img, background_mask_3ch)

        # 提取并应用增强后的文字
        enhanced_text = cv2.bitwise_and(text_enhanced, text_mask_3ch)

        # 合并背景和增强文字
        result = cv2.add(background, enhanced_text)

        # 转换为PIL格式进行最终调整
        result_pil = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))

        # 应用锐化滤镜
        result_pil = result_pil.filter(ImageFilter.UnsharpMask(radius=2, percent=150, threshold=3))

        # 增加对比度
        enhancer = ImageEnhance.Contrast(result_pil)
        result_pil = enhancer.enhance(1.8)

        # 锐化文字边缘
        enhancer = ImageEnhance.Sharpness(result_pil)
        result_pil = enhancer.enhance(3.0)

        # 如果原图有alpha通道，添加回去
        if has_alpha:
            alpha = original_img.split()[-1]
            result_pil = result_pil.convert("RGBA")
            result_pil.putalpha(alpha)

        # 生成输出路径
        directory, filename = os.path.split(image_path)
        name, ext = os.path.splitext(filename)
        output_path = os.path.join(directory, f"{name}_enhanced{ext}")

        # 保存处理后的图片
        result_pil.save(output_path, dpi=(300, 300))
        print(f"增强后的图片已保存至: {output_path}")
        return output_path

    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        raise


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        print("使用方法: python enhance_text_clarity.py <图片路径>")
        sys.exit(1)

    input_image = sys.argv[1]
    if not os.path.isfile(input_image):
        print(f"错误: 文件 '{input_image}' 不存在")
        sys.exit(1)

    try:
        enhance_image_clarity(input_image)
    except Exception as e:
        print(f"处理失败: {str(e)}")
        sys.exit(1)